from __future__ import annotations
import numpy as np
import manim as mnm
from typing import Sequence
from misc.character_design import *


MY_FONT = 'Malgun Gothic Semilight'


def get_series_axes(value_series: np.ndarray, title, height, year=False) -> tuple[mnm.Axes, mnm.VGroup]:
    """

    :rtype: tuple of axes and group of axes title and tips
    """
    assert isinstance(value_series, np.ndarray)
    x_step = 250 if year else 30
    axes = mnm.Axes(x_range=[0, value_series.shape[0], x_step],
                    y_range=[value_series.min(), value_series.max(),
                             (value_series.max() - value_series.min()) / 10],
                    y_length=height,
                    # x_axis_config={"include_ticks": False},
                    y_axis_config={"font_size": 18, "include_numbers": True,
                                   "decimal_number_config": {"num_decimal_places": 0}},
                    tips=False
                    )
    # Labels for the x-axis and y-axis.
    # month_labels = mnm.VGroup()
    # note: use legend not label
    # month_labels.add(axes.get_y_axis_label(mnm.Text("净值", font_size=18)))
    # month_labels.add(axes.get_x_axis_label(mnm.Text("日期", font_size=18), edge=mnm.DOWN, direction=mnm.DOWN,
    #                                        buff=mnm.MED_LARGE_BUFF))
    title = mnm.Text(title, font_size=18, font=MY_FONT).next_to(axes.get_top(), direction=mnm.DOWN)
    month_group = mnm.VGroup(axes, title)
    return axes, month_group


class TimeSeriesAxes:
    def __init__(self, value_series: np.ndarray, title, height, year=False):
        self.value_series = value_series
        self.sample_length = value_series.shape[0]
        self.axes, self.axes_group = get_series_axes(value_series, title, height, year)
        self.dot = None
        self.trace = None
        self.wandering_lines = None

    def __getattr__(self, item):
        if item in self.axes.__dict__:
            return getattr(self.axes, item)
        else:
            raise AttributeError('Unknown attribute : %s' % item)

    def add(self, *mob):
        self.axes_group.add(*mob)

    def add_date_ticks(self, date_labels: dict, **kwargs):
        self.axes.x_axis.add_labels(date_labels, font_size=18, buff=mnm.SMALL_BUFF, **kwargs)

    def coords_to_point(self, *coords):
        return self.axes.coords_to_point(*coords)

    def index_to_point(self, index):
        return self.coords_to_point(index, self.value_series[index])

    def get_series(self, index):
        return self.value_series[index]

    def generate_wandering_line(self, st, ed, color):
        wandering_lines = mnm.VMobject(color=color, stroke_width=1)
        zero_point = self.index_to_point(st)
        wandering_lines.set_points_as_corners([zero_point, zero_point])
        for idx in range(st+1, ed):
            wandering_lines.add_line_to(self.index_to_point(idx))
        return wandering_lines

    def wandering_series(self, dot_color=mnm.GREY) -> tuple[mnm.VGroup, mnm.Animation]:
        # cp_color = mnm.GREY
        zero_point = self.index_to_point(0)
        self.dot = mnm.Dot(zero_point, color=dot_color)
        self.trace = mnm.TracedPath(self.dot.get_center, color=dot_color)
        self.wandering_lines = self.generate_wandering_line(0, self.sample_length, dot_color)
        self.axes_group.add(self.dot, self.trace)
        return self.axes_group, mnm.MoveAlongPath(self.dot, self.wandering_lines)

    def clear_dot(self):
        self.axes_group -= self.dot
        self.axes_group -= self.trace
        self.axes_group += self.wandering_lines


class Outline(mnm.Scene):
    def construct(self):
        title = mnm.Text("情绪与长期持有", color=mnm.WHITE)
        title.to_edge(mnm.UP)
        title.shift(4 * mnm.UP)
        line = mnm.Line(mnm.config.frame_x_radius * mnm.UP, mnm.config.frame_x_radius * mnm.RIGHT)
        line.shift(np.array([0, title.get_bottom()[1] - 0.2, 0]))
        self.add(title, line)
        t1 = mnm.Text("1. 闭上眼睛", color=mnm.WHITE)
        t2 = mnm.Text("2. 转移注意", color=mnm.WHITE)
        x = mnm.VGroup(t1, t2).arrange(direction=mnm.DOWN, aligned_edge=mnm.UP).scale(0.8).next_to(
            title.get_left() + 2 * mnm.DOWN)
        x.set_opacity(0.5)
        x.submobjects[1].set_opacity(1)
        self.add(x)
        author = mnm.Text("@隆隆.ago", color=mnm.WHITE).scale(0.5).shift(mnm.frame.right_side + mnm.frame.bottom + mnm.UL)
        self.add(author)
        self.play(mnm.FadeIn(x))
        self.wait(2)
        self.play(mnm.FadeOut(*self.mobjects))


class OpeningQuote(mnm.Scene):
    def construct(self, quote='“幸福是在向山顶努力攀登过程中的种种经历和感受”', author="有知有行",
                  fade_in_kwargs={"lag_ratio": 0.5, "rate_func": mnm.linear, "run_time": 2},
                  text_size=48, use_quotation_marks=False, top_buff=1.0, author_buff=1.0, wait_time=2):
        self.quote = self.get_quote(quote)
        self.author = self.get_author(author)
        self.fade_in_kwargs = fade_in_kwargs
        self.text_size = text_size
        self.use_quotation_marks = use_quotation_marks
        self.highlighted_quote_terms = {"幸福": mnm.BLUE}
        self.quote_arg_separator = " "

        self.top_buff = top_buff
        self.author_buff = author_buff

        self.play(mnm.FadeIn(self.quote, **self.fade_in_kwargs))
        self.wait(wait_time)
        self.play(mnm.Write(self.author, run_time=1))
        self.wait()

    def get_quote(self, quote, max_width=mnm.config.frame_width - 1):
        if isinstance(quote, str):
            if self.use_quotation_marks:
                quote = mnm.Text("``%s''" % quote.strip(), font_size=self.text_size)
            else:
                quote = mnm.Text(r"%s" % quote.strip(), font_size=self.text_size)
        else:
            if self.use_quotation_marks:
                words = [self.text_size + " ``"] + list(quote) + ["''"]
            else:
                words = [self.text_size] + list(quote)
            print(words)
            quote = mnm.Text(*words, font_size=self.text_size)
            # TODO, make less hacky
            if self.quote_arg_separator == " ":
                quote[0].shift(0.2 * mnm.RIGHT)
                quote[-1].shift(0.2 * mnm.LEFT)
        for term, color in self.highlighted_quote_terms.items():
            quote.set_color_by_tex(term, color)
        quote.to_edge(mnm.UP, buff=self.top_buff)
        if quote.get_width() > max_width:
            quote.set_width(max_width)
        return quote

    def get_author(self, author):
        author = mnm.Text(" ————" + author, font_size=self.text_size)
        author.next_to(self.quote, mnm.DOWN, buff=self.author_buff)
        author.set_color(mnm.YELLOW)
        return author


class MyMovieRef(mnm.Scene):
    def construct(self):
        self.create_title_and_yan()

    def create_title_and_yan(self):
        self.title = mnm.Text("Reference")
        self.title.scale(1.5)
        self.title.to_edge(mnm.UP, buff=mnm.MED_SMALL_BUFF)
        from misc.character_design.pi_creature import Yan
        randy, morty = self.pi_creatures = mnm.VGroup(Yan(), Yan(color=mnm.GRAY))
        for pi, vect in (randy, mnm.LEFT), (morty, mnm.RIGHT):
            pi.set(height=self.title.height)
            # pi.change_mode("thinking")
            pi.look_at(mnm.DOWN)
            pi.next_to(self.title, vect, buff=mnm.MED_LARGE_BUFF)
        line = mnm.DashedLine(mnm.config.frame_x_radius * mnm.LEFT, mnm.config.frame_x_radius * mnm.RIGHT)
        line.next_to(self.title, direction=mnm.DOWN, buff=mnm.MED_SMALL_BUFF)
        self.add_foreground_mobjects(self.title, randy, morty, line)


if __name__ == "__main__":
    with mnm.tempconfig({"quality": "low_quality", "preview": True, "frame_rate": 50}):
        print(mnm.config["frame_rate"])
        scene = MyMovieRef()
        scene.render()
    pass
